import json
from collections.abc import Callable, Iterator, Mapping
from typing import TypeVar

from zkl_pyutils_fsspec import Fs, FsLike, resolve_child_fs, resolve_fs
from zkl_pyutils_serialization import dump_json_value, parse_json_value

from .loader import RecordJson, load_record, record_json_file_name

AnyWriter = TypeVar('AnyWriter')
AnyReader = TypeVar('AnyReader')


class DictRecordReader(Mapping[str, AnyReader]):
    def __init__(self, fs: Fs, *, children_path: dict[str, str]):
        self._fs = fs
        self._children_path = children_path
        self._children_reader = {}

    def refresh(self, clean: bool):
        with self._fs.open(record_json_file_name, 'rt', encoding='utf-8') as fp:
            record_json = parse_json_value(json.load(fp), RecordJson)
        self._children_path = record_json.kwargs['children_path']
        if clean:
            self._children_reader.clear()

    def __len__(self) -> int:
        return len(self._children_path)

    def __iter__(self) -> Iterator[str]:
        return iter(self._children_path)

    def __getitem__(self, key: str, /) -> AnyReader:
        reader = self._children_reader.get(key)
        if reader is None:
            path = self._children_path[key]
            fs = resolve_child_fs(self._fs, path)
            reader = load_record(fs)
            self._children_reader[key] = reader
        return reader


class DictRecordWriter:
    def __init__(self, fs: FsLike, *, reader_loader: Callable = DictRecordReader):
        fs = resolve_fs(fs)
        fs.makedirs("", exist_ok=True)

        self._fs = fs
        self._reader_loader = reader_loader

        if fs.exists(record_json_file_name):
            with fs.open(record_json_file_name, 'rt', encoding='utf-8') as fp:
                record_json = parse_json_value(json.load(fp), RecordJson)
            self._children_path = record_json.kwargs['children_path']
        else:
            self._children_path = {}
            self._save_record_json()
        self._children_writer = {}

    def _save_record_json(self):
        record_json = RecordJson.of(self._reader_loader, children_path=self._children_path)
        with self._fs.open(record_json_file_name, 'wt', encoding='utf-8') as fp:
            json.dump(dump_json_value(record_json), fp, indent=2, ensure_ascii=False)

    def get(self, key: str, writer_factory: Callable[[Fs], AnyWriter] | None = None) -> AnyWriter:
        writer = self._children_writer.get(key)
        if writer is None:
            if writer_factory is None:
                raise KeyError
            path = self._children_path.get(key)
            if path is None:
                path = _key_to_path(key)
                self._children_path[key] = path
                self._save_record_json()
            fs = resolve_child_fs(self._fs, path)
            writer = writer_factory(fs)
            self._children_writer[key] = writer
        return writer


def _key_to_path(key: str) -> str:
    key = key.replace(':', '_')
    key = key.replace('\\', '_')
    key = key.replace('?', '_')
    key = key.replace('*', '_')
    key = key.replace('"', '_')
    key = key.replace('<', '_')
    key = key.replace('>', '_')
    key = key.replace('|', '_')
    key = key.replace('%', '_')
    key = key.replace('&', '_')
    return key
